package com.cfp4cloud.cfp.knowledge.support.provider;

import cn.hutool.extra.spring.SpringUtil;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.cfp4cloud.cfp.common.core.constant.enums.YesNoEnum;
import com.cfp4cloud.cfp.knowledge.config.properties.AiKnowledgeProperties;
import com.cfp4cloud.cfp.knowledge.entity.AiModelEntity;
import com.cfp4cloud.cfp.knowledge.mapper.AiModelMapper;
import com.cfp4cloud.cfp.knowledge.support.constant.ModelTypeEnums;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.internal.Utils;
import dev.langchain4j.model.embedding.DimensionAwareEmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.filter.Filter;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import lombok.experimental.UtilityClass;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.Nullable;

import java.util.Objects;

/**
 * 内存向量存储提供者
 * <p>
 * 提供基于内存的向量存储和检索功能，用于AI随航等轻量级场景 支持文本向量化、存储和相似度搜索，数据存储在内存中 TODO: 后续需要实现持久化功能，支持系统重启后的数据恢复
 *
 * @author chenda
 * @date 2025/03/21
 */
@Slf4j
@UtilityClass
public class MemoryEmbeddingProvider {

	/**
	 * 内存向量存储实例 注意：数据仅存储在内存中，系统重启后会丢失
	 */
	public static final InMemoryEmbeddingStore<TextSegment> EMBEDDING_STORE = new InMemoryEmbeddingStore<>();

	/**
	 * 临时ID常量，用于临时存储场景
	 */
	public static final String TEMP_ID = "temp_id";

	/**
	 * 添加文本段到向量存储
	 * <p>
	 * 自动生成随机ID并进行向量化存储
	 * @param textSegment 要存储的文本段
	 */
	public void add(TextSegment textSegment) {
		add(Utils.randomUUID(), textSegment);
	}

	/**
	 * 添加文本段到向量存储（指定ID）
	 * <p>
	 * 将文本向量化后存储，支持后续的相似度搜索
	 * @param id 文档ID
	 * @param textSegment 要存储的文本段
	 */
	public void add(String id, TextSegment textSegment) {
		DimensionAwareEmbeddingModel embeddingModel = getEmbeddingModel();
		if (embeddingModel == null)
			return;

		// 文本向量化
		dev.langchain4j.model.output.Response<Embedding> embeddingResponse = embeddingModel.embed(textSegment.text());
		Embedding embedding = embeddingResponse.content();
		// 存储向量和原始文本
		EMBEDDING_STORE.add(id, embedding, textSegment);
	}

	/**
	 * 获取默认的嵌入模型
	 * <p>
	 * 从数据库配置中获取默认的向量模型
	 * @return 嵌入模型实例，如果没有配置则返回null
	 */
	@Nullable
	private static DimensionAwareEmbeddingModel getEmbeddingModel() {
		AiModelMapper modelMapper = SpringUtil.getBean(AiModelMapper.class);

		AiModelEntity aiModelEntity = modelMapper.selectOne(Wrappers.<AiModelEntity>lambdaQuery()
			.eq(AiModelEntity::getDefaultModel, YesNoEnum.YES.getCode())
			.eq(AiModelEntity::getModelType, ModelTypeEnums.EMBEDDING.getType()), false);

		if (Objects.isNull(aiModelEntity)) {
			log.warn("没有默认的向量模型，请检查模型配置，AI随航无法使用");
			return null;
		}

		ModelProvider modelProvider = SpringUtil.getBean(ModelProvider.class);
		return modelProvider.getEmbeddingModel(aiModelEntity.getName());
	}

	/**
	 * 相似度搜索（使用默认配置）
	 * <p>
	 * 使用系统默认的topK和最小分数进行搜索
	 * @param text 搜索文本
	 * @param filters 过滤条件
	 * @return 搜索结果
	 */
	public EmbeddingSearchResult<TextSegment> search(String text, Filter... filters) {
		AiKnowledgeProperties properties = SpringUtil.getBean(AiKnowledgeProperties.class);
		return search(text, properties.getInMemorySearch().getTopK(), properties.getInMemorySearch().getMinScore(),
				filters);
	}

	/**
	 * 相似度搜索（自定义参数）
	 * <p>
	 * 将文本向量化后在向量库中进行相似度搜索
	 * @param text 搜索文本
	 * @param topK 返回结果数量
	 * @param minScore 最低相似度分数
	 * @param filters 过滤条件
	 * @return 搜索结果
	 */
	public EmbeddingSearchResult<TextSegment> search(String text, int topK, double minScore, Filter... filters) {
		// 获取嵌入模型
		DimensionAwareEmbeddingModel embeddingModel = getEmbeddingModel();
		if (embeddingModel == null)
			return null;

		// 搜索文本向量化
		Embedding embedding = getEmbeddingModel().embed(text).content();

		// 构建搜索请求
		EmbeddingSearchRequest.EmbeddingSearchRequestBuilder builder = EmbeddingSearchRequest.builder()
			.queryEmbedding(embedding)
			.minScore(minScore)
			.maxResults(topK);

		// 添加过滤条件
		for (Filter filter : filters) {
			builder.filter(filter);
		}

		// 执行向量搜索
		return EMBEDDING_STORE.search(builder.build());
	}

}
